package main

import (
	"context"
	"flag"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/containerd/containerd/log"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"

	"onyx-snapshotter/pkg/config"
	"onyx-snapshotter/pkg/metrics"
	"onyx-snapshotter/pkg/snapshotter"

	snapshotsapi "github.com/containerd/containerd/api/services/snapshots/v1"
)

var (
	version = "v2.0.0"

	configPath  = flag.String("config", "/etc/onyx/config.toml", "path to config file")
	logLevel    = flag.String("log-level", "info", "log level (debug, info, warn, error)")
	address     = flag.String("address", "/run/containerd/onyx-snapshotter.sock", "gRPC socket address")
	metricsAddr = flag.String("metrics-address", ":9100", "Prometheus metrics address")
	showVersion = flag.Bool("version", false, "show version and exit")
)

func main() {
	flag.Parse()

	if *showVersion {
		logrus.Infof("onyx-snapshotter version %s", version)
		return
	}

	// Setup logging
	level, err := logrus.ParseLevel(*logLevel)
	if err != nil {
		logrus.WithError(err).Fatal("invalid log level")
	}
	logrus.SetLevel(level)
	logrus.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: "2006-01-02T15:04:05.000Z07:00",
		FieldMap: logrus.FieldMap{
			logrus.FieldKeyTime:  "timestamp",
			logrus.FieldKeyLevel: "level",
			logrus.FieldKeyMsg:   "message",
		},
	})

	ctx := log.WithLogger(context.Background(), log.L)
	log.G(ctx).WithField("version", version).Info("starting onyx-snapshotter")

	// Load configuration
	cfg, err := config.Load(*configPath)
	if err != nil {
		log.G(ctx).WithError(err).Fatal("failed to load config")
	}

	// Initialize metrics
	if err := metrics.Initialize(); err != nil {
		log.G(ctx).WithError(err).Fatal("failed to initialize metrics")
	}

	// Start metrics server
	go func() {
		http.Handle("/metrics", promhttp.Handler())
		http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("OK"))
		})
		log.G(ctx).WithField("address", *metricsAddr).Info("starting metrics server")
		if err := http.ListenAndServe(*metricsAddr, nil); err != nil {
			log.G(ctx).WithError(err).Error("metrics server failed")
		}
	}()

	// Create snapshotter
	sn, err := snapshotter.New(ctx, cfg)
	if err != nil {
		log.G(ctx).WithError(err).Fatal("failed to create snapshotter")
	}
	defer sn.Close()

	// Remove existing socket
	if err := os.RemoveAll(*address); err != nil {
		log.G(ctx).WithError(err).Fatal("failed to remove existing socket")
	}

	// Create gRPC server with options
	opts := []grpc.ServerOption{
		grpc.MaxRecvMsgSize(16 * 1024 * 1024), // 16MB
		grpc.MaxSendMsgSize(16 * 1024 * 1024),
	}
	rpc := grpc.NewServer(opts...)
	service := snapshotter.NewService(sn)
	snapshotsapi.RegisterSnapshotsServer(rpc, service)

	// Listen on Unix socket
	listener, err := net.Listen("unix", *address)
	if err != nil {
		log.G(ctx).WithError(err).Fatal("failed to listen")
	}

	// Set socket permissions
	if err := os.Chmod(*address, 0o660); err != nil {
		log.G(ctx).WithError(err).Warn("failed to set socket permissions")
	}

	// Handle shutdown signals
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		sig := <-sigCh
		log.G(ctx).WithField("signal", sig).Info("received shutdown signal")
		rpc.GracefulStop()
		listener.Close()
	}()

	log.G(ctx).WithFields(map[string]interface{}{
		"address": *address,
		"version": version,
		"metrics": *metricsAddr,
	}).Info("onyx-snapshotter started successfully")

	metrics.ServiceStartTime.SetToCurrentTime()

	if err := rpc.Serve(listener); err != nil {
		log.G(ctx).WithError(err).Fatal("failed to serve")
	}
}
